import requests
from PIL import Image, ImageOps
import cv2
import os
import random

# 定义下载图片并保存到本地的函数
def download_image(url, save_path):
    response = requests.get(url)
    if response.status_code == 200:
        with open(save_path, 'wb') as f:
            f.write(response.content)
        return True
    return False

# 定义数据增强的函数
def augment_image(image):
    # 随机选择一种增强方式
    augmentations = [ImageOps.mirror, ImageOps.flip, ImageOps.rotate]
    augmented_image = image.copy()
    chosen_aug = random.choice(augmentations)
    if chosen_aug == ImageOps.mirror:
        augmented_image = ImageOps.mirror(augmented_image)
    elif chosen_aug == ImageOps.flip:
        augmented_image = ImageOps.flip(augmented_image)
    elif chosen_aug == ImageOps.rotate:
        augmented_image = ImageOps.rotate(augmented_image, random.randint(-45, 45))
    return augmented_image

# 主函数
def main():
    input_file = 'face.txt'
    output_dir = 'augmented_images'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    with open(input_file, 'r') as f:
        urls = f.readlines()

    # 下载并保存原始图片
    for i, url in enumerate(urls[:50]):
        url = url.strip()
        image_path = os.path.join(output_dir, f'original_{i}.jpg')
        if download_image(url, image_path):
            print(f"Downloaded {url} to {image_path}")

    # 对每张原始图片进行数据增强，生成额外的图片
    for i in range(len(urls)):
        url = urls[i].strip()
        image_path = os.path.join(output_dir, f'original_{i}.jpg')
        original_image = Image.open(image_path)

        for j in range(4):  # 每张图片生成4种增强版本
            augmented_image = augment_image(original_image)
            new_image_path = os.path.join(output_dir, f'augmented_{i}_{j}.jpg')
            augmented_image.save(new_image_path)
            print(f"Saved augmented image to {new_image_path}")

if __name__ == '__main__':
    main()